import win32gui
import win32ui
import win32con
import numpy as np
import cv2
from typing import Tuple, Union, Optional
from enum import Enum
from ...window.window_manager import WindowInfo
from ...window.draw_utils import WindowDraw, DrawColor

class CaptureAnchor(Enum):
    """截图锚点位置"""
    LEFT = "left"
    RIGHT = "right"
    TOP = "top"
    BOTTOM = "bottom"
    LEFT_TOP = "left_top"
    RIGHT_TOP = "right_top"
    LEFT_BOTTOM = "left_bottom"
    RIGHT_BOTTOM = "right_bottom"
    CENTER = "center"

def _convert_to_pixels(value: Union[int, float, str], base_value: int) -> int:
    """将值转换为像素值"""
    if isinstance(value, str) and value.endswith('%'):
        return int(float(value.rstrip('%')) * base_value / 100)
    elif isinstance(value, float) and 0 <= value <= 1:
        return int(value * base_value)
    return int(value)

def capture_window(
    window_info: WindowInfo,
    start_point: Tuple[Union[int, float, str], Union[int, float, str]] = (0, 0),
    end_point: Optional[Tuple[Union[int, float, str], Union[int, float, str]]] = None,
    size: Optional[Tuple[Union[int, float, str], Union[int, float, str]]] = None,
    anchor: Union[str, CaptureAnchor] = CaptureAnchor.LEFT_TOP,
    debug: bool = False
) -> Optional[np.ndarray]:
    """截取窗口指定区域"""
    try:
        # 获取窗口大小
        window_width = window_info.size[0]
        window_height = window_info.size[1]
        
        # 转换起始点坐标
        start_x = _convert_to_pixels(start_point[0], window_width)
        start_y = _convert_to_pixels(start_point[1], window_height)
        
        # 计算宽度和高度
        if end_point is not None:
            # 使用终点计算
            end_x = _convert_to_pixels(end_point[0], window_width)
            end_y = _convert_to_pixels(end_point[1], window_height)
            width = end_x - start_x
            height = end_y - start_y
        elif size is not None:
            # 使用大小计算
            width = _convert_to_pixels(size[0], window_width)
            height = _convert_to_pixels(size[1], window_height)
        else:
            # 默认使用整个窗口
            width = window_width - start_x
            height = window_height - start_y
        
        # 根据锚点调整坐标
        if isinstance(anchor, str):
            anchor = CaptureAnchor(anchor.lower())
            
        if anchor == CaptureAnchor.RIGHT:
            start_x = window_width - width - start_x
        elif anchor == CaptureAnchor.BOTTOM:
            start_y = window_height - height - start_y
        elif anchor == CaptureAnchor.RIGHT_BOTTOM:
            start_x = window_width - width - start_x
            start_y = window_height - height - start_y
        elif anchor == CaptureAnchor.CENTER:
            start_x = (window_width - width) // 2 + start_x
            start_y = (window_height - height) // 2 + start_y
        
        # 如果开启调试模式，绘制调试框
        if debug:
            WindowDraw.draw_rect(
                window_info.hwnd,
                start_x,
                start_y,
                width,
                height,
                color=DrawColor.RED,
                duration_ms=500
            )
        
        # 创建设备上下文
        hwndDC = win32gui.GetWindowDC(window_info.hwnd)
        mfcDC = win32ui.CreateDCFromHandle(hwndDC)
        saveDC = mfcDC.CreateCompatibleDC()
        
        # 创建位图
        saveBitMap = win32ui.CreateBitmap()
        saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
        saveDC.SelectObject(saveBitMap)
        
        # 复制图像
        saveDC.BitBlt((0, 0), (width, height), mfcDC, (start_x, start_y), win32con.SRCCOPY)
        
        # 转换为numpy数组
        bmpinfo = saveBitMap.GetInfo()
        bmpstr = saveBitMap.GetBitmapBits(True)
        img = np.frombuffer(bmpstr, dtype='uint8')
        img.shape = (height, width, 4)
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        
        # 清理资源
        win32gui.DeleteObject(saveBitMap.GetHandle())
        saveDC.DeleteDC()
        mfcDC.DeleteDC()
        win32gui.ReleaseDC(window_info.hwnd, hwndDC)
        
        return img
        
    except Exception as e:
        print(f"截图失败: {e}")
        return None 